import pybuda._C.backend_api
import pybuda._C.balancer
import torch
from typing import Dict, List, Optional

class CompileRequest:
    def __init__(self, netlist_path: str, output_dir: str, backend_config: pybuda._C.backend_api.BackendConfig, inputs: List[PyBudaTensorDesc], input_runtime_transforms: List[str], constants: List[PyBudaTensorDesc], parameters: List[PyBudaTensorDesc], outputs: List[PyBudaTensorDesc], output_runtime_transforms: List[str]) -> None: ...

class Program:
    def __init__(self, name: str, params: Dict[str, str]) -> None: ...

class PyBudaTensorDesc:
    def __init__(self, name: str, shape: List[int], ptr: int = ..., constant: Optional[torch.Tensor] = ...) -> None: ...
    @property
    def constant(self) -> Optional[torch.Tensor]: ...
    @property
    def name(self) -> str: ...
    @property
    def ptr(self) -> int: ...
    @property
    def shape(self) -> List[int]: ...

class TTDevice:
    def __init__(self, *args, **kwargs) -> None: ...
    def compile(self, arg0: CompileRequest) -> Workload: ...
    def dispatch(self, arg0: Workload, arg1: List[Program], arg2: List[torch.Tensor], arg3: Dict[str, pybuda._C.balancer.OutputHostTM]) -> List[torch.Tensor]: ...
    def str(self) -> str: ...
    def torch_device(self) -> torch.device: ...
    @property
    def arch(self) -> pybuda._C.backend_api.BackendDevice: ...
    @property
    def cluster_yaml(self) -> str: ...
    @property
    def index(self) -> int: ...
    @property
    def mmio(self) -> bool: ...
    @property
    def soc_desc_yaml(self) -> str: ...
    @property
    def type(self) -> pybuda._C.backend_api.BackendType: ...

class Workload:
    def __init__(self, *args, **kwargs) -> None: ...
    @property
    def backend(self) -> pybuda._C.backend_api.BackendApi: ...
    @property
    def constants(self) -> List[PyBudaTensorDesc]: ...
    @property
    def inputs(self) -> List[PyBudaTensorDesc]: ...
    @property
    def outputs(self) -> List[PyBudaTensorDesc]: ...
    @property
    def parameters(self) -> List[PyBudaTensorDesc]: ...

def get_available_devices(*args, **kwargs): ...
def get_default_device(*args, **kwargs): ...
def is_created_on_device(arg0: torch.Tensor) -> bool: ...
def original_shape(arg0: torch.Tensor) -> List[int]: ...
def push_tensor(arg0: pybuda._C.backend_api.BackendApi, arg1: PyBudaTensorDesc, arg2: torch.Tensor, arg3: str) -> None: ...
